import { type Rank } from '@tensorflow/tfjs'
import tf from '../tf'
import { LoadStatus, TorchModeJson } from '../types'
import { BaseConvolutionModel } from './BaseConvolutionModel'
import BaseLoadHandler from './loadHandler/BaseLoadHandler'
import SimpleLoadHandler from './loadHandler/SimpleLoadHandler'
import { upscaleModelUrls } from './urls'

const getLayers = (json: TorchModeJson[]) => {
  const layers: any[] = []
  layers.push(
    tf.layers.inputLayer({
      inputShape: [null, null, 3],
    })
  )
  json.forEach((item) => {
    if (item.class_name != 'nn.SpatialConvolutionMM') {
      throw new Error('Not implemented')
    }
    layers.push(
      tf.layers.conv2d({
        strides: [item.dW, item.dH],
        kernelSize: [item.kW, item.kH],
        filters: item.nOutputPlane,
        padding: 'valid',
        inputDType: 'float32',
        useBias: true,
        activation: 'linear',
      }),
      tf.layers.leakyReLU({
        alpha: 0.1,
      })
    )
  })
  return layers
}
function getWeightByItem(
  json: TorchModeJson[],
  item: TorchModeJson,
  index: number
) {
  let channel = 3
  if (index != 0) {
    channel = json[index - 1].nOutputPlane
  }
  return tf.tidy(() => {
    return tf
      .tensor(
        item.weight,
        [item.nOutputPlane, channel, item.kW, item.kH],
        'float32'
      )
      .transpose([2, 3, 1, 0])
  })
}
const getWeights = (json: TorchModeJson[]): tf.Tensor<Rank>[] => {
  // 定义一个空数组，用于存储权重
  const weights: tf.Tensor<Rank>[] = []
  json.forEach((item, index) => {
    if (item.class_name == 'nn.SpatialConvolutionMM') {
      weights.push(getWeightByItem(json, item, index))
      weights.push(tf.tensor(item.bias, [item.nOutputPlane], 'float32'))
    }
  })
  return weights
}

class ScaleConvModel extends BaseConvolutionModel {
  private model?: tf.Sequential
  private loadHander: BaseLoadHandler
  constructor() {
    super()
    this.loadHander = new SimpleLoadHandler(
      'ScaleConvModel',
      upscaleModelUrls.waifu_2x,
      1024 * 1024 * 2.6,
      (data) => {
        const model = tf.sequential({
          name: 'waifu2x_scale2x',
          layers: getLayers(data),
        })
        model.setWeights(getWeights(data))
        this.model = model
      }
    )
    // this.loadHander
    //   .supportCacheLoad()
    //   .then((falg) => falg && this.loadHander.load())
  }
  public getLoadHander(): BaseLoadHandler {
    return this.loadHander
  }
  //卷积了7层，卷积一层padding-1
  public getInputPadding(): number {
    return 7
  }

  public predictSize(width: number, height: number) {
    width *= 2
    height *= 2
    return { width, height }
  }
  public initResize(width: number, height: number) {
    width *= 2
    height *= 2
    return { width, height }
  }

  public initInputTensor(input: tf.Tensor3D): tf.Tensor3D {
    input = super.initInputTensor(input)
    const out = tf.image.resizeBilinear(
      input,
      [input.shape[0] * 2, input.shape[1] * 2],
      true
    )
    input.dispose()
    return out
  }
  public predict(input: tf.Tensor3D): tf.Tensor3D {
    if (!this.isLoaded) {
      const status = this.getLoadHander().getStatus()
      if (status === LoadStatus.loading) {
        throw new Error('模型正在加载中，请稍后')
      } else if (status === LoadStatus.error) {
        throw new Error('模型加载失败，请检查网络连接，或者尝试重新加载模型')
      } else if (status === LoadStatus.init) {
        throw new Error('模型未加载，请先加载模型')
      } else {
        throw new Error('model no ready')
      }
    }
    const x = input.reshape([1, ...input.shape])
    const y = this.model!.predict(x) as tf.Tensor4D
    const output = y.as3D(y.shape[1], y.shape[2], y.shape[3])
    tf.dispose([x, y])
    return output
  }
}

export default new ScaleConvModel()
